import redis

class URLManagementMiddleware:
    def __init__(self, redis_host='localhost', redis_port=6379, db=0):
        self.redis = redis.StrictRedis(host=redis_host, port=redis_port, db=db)
        self.to_crawl_url_set_key = 'to_crawl_urls'

    def add_url(self, url, priority=1):
        """
        添加待爬取 URL 到 Redis 有序集合中，根据优先级设置分数
        """
        self.redis.zadd(self.to_crawl_url_set_key, {url: priority})

    def get_next_url(self):
        """
        获取下一个待爬取 URL，即分数最高的 URL
        """
        result = self.redis.zpopmax(self.to_crawl_url_set_key)
        if result:
            return result[0][0].decode('utf-8')
        return None
    
    